import importlib
import inspect
import math
import re
import threading
import time
import traceback
import types
from functools import wraps
import ast
import json
import os
import sys
from typing import Union

from pypinyin import pinyin, Style
import one_runtime
import util

# 脚本目录
script_directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), "scripts")

# 脚本线程局部变量
script_thread_local = threading.local()

# 合并大模型花费线程局部变量
merge_thread_local = threading.local()

# 脚本缓存
__script_schemas: dict[str, dict] = {}

# 脚本缓存
__script_modules: dict = {}


def get_script_schemas():
    return __script_schemas


def get_script_modules():
    return __script_modules


# 限流信息
__function_buckets = {}

# 监控信息
__function_monitors = {}


# 任务上下文
class NodeTaskCtx:
    def __init__(self, task_name: str = None, task_source: dict = None, script_name: str = None, node: str = None, data_group: str = None, req_id: str = None, plan_config: dict = None, test_name: str = None, debug_param: str = None, pass_through: dict = None, prepare_result: dict = None):
        '''
        :param task_name: 任务名称
        :param task_source: 任务源配置
        :param script_name: 任务名称
        :param node: 节点
        :param data_group:数据分组
        :param req_id: 请求id
        :param plan_config: 方案配置 默认填充该节点在线的配置
        :param test_name: 实验名称
        :param debug_param: 调试参数 一定不要强依赖任务配置的传参，建议作为兜底/调试参数
        :param pass_through: 程序透传参数  层层传递 应和节点无关的全局属性 例如 线上/预发
        :param prepare_result: 准备结果
        '''
        self.task_name = task_name
        self.task_source = task_source
        self.script_name = script_name
        self.node = node
        self.data_group = data_group
        self.req_id = req_id
        self.plan_config = plan_config
        self.test_name = test_name
        self.debug_param = debug_param
        self.pass_through = pass_through if isinstance(pass_through, dict) else {}
        self.prepare_result = prepare_result if isinstance(prepare_result, dict) else {}

    def debug_param_as_list(self):
        # 配置使用字符串拆分成列表
        if not self.debug_param:
            return None
        return self.debug_param.strip().split(',')

    def debug_param_as_dict(self):
        # 配置使用字符串拆分成字典
        try:
            if not self.debug_param:
                return None
            return json.loads(self.debug_param.strip())
        except Exception as e:
            raise ValueError('调试参数解析失败,预期为json。' + str(e))

    # 获取方案配置
    def get_plan_config(self, key, def_val=None):
        # 配置使用字符串拆分成字典
        if not self.plan_config:
            return def_val
        return self.plan_config.get(key, def_val)

    # 获取透传参数
    def get_pass_through(self, key, def_val=None):
        # 配置使用字符串拆分成字典
        if not self.pass_through:
            return def_val
        return self.pass_through.get(key, def_val)

    def to_dict(self):
        return {
            "task_name": self.task_name,
            "task_source": self.task_source,
            "script_name": self.script_name,
            "plan_config": self.plan_config,
            "node": self.node,
            "data_group": self.data_group,
            "req_id": self.req_id,
            "test_name": self.test_name,
            'debug_param': self.debug_param,
            'pass_through': self.pass_through
        }

    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False,default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()


# 节点运行结果
class NodeResult:
    def __init__(self, content_key=None, content_type=None, content=None, ext_info: dict = None, status=None):
        '''
        :param content_key:  这条数的业务key是什么 框架不使用
        :param content_type: 数据类型 默认是default
        :param content: 数据
        :param ext_info:扩展数据 字典
        :param status: 状态码 框架不使用
        '''
        self.content_key = content_key
        try:
            self.content = json.loads(content)
        except:
            self.content = content
        self.content_type = content_type
        if not self.content_type:
            self.content_type = 'default'
        self.ext_info = ext_info if ext_info else {}
        if not isinstance(self.ext_info, dict):
            try:
                self.ext_info = json.loads(str(self.ext_info))
            except:
                raise ValueError('NodeResult.ext_info must be dict')
        if status != None:
            self.ext_info.update({
                'status': status
            })

    def to_dict(self):
        return {
            "content_key": self.content_key,
            "content_type": self.content_type,
            "content": self.content,
            "ext_info": self.ext_info
        }

    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False, default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()


class EvalData:
    def __init__(self, id: int = None, data_group: str = None, content_type: str = None, content: Union[str, dict] = None, ext_info: dict = None):
        self.id = id
        self.data_group = data_group
        self.content_type = content_type
        self.content = content
        self.ext_info = ext_info

    def to_dict(self):
        return {
            'id': self.id,
            "data_group": self.data_group,
            "content_type": self.content_type,
            "content": self.content,
            "ext_info": self.ext_info
        }

    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False,default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()


# 评测数据眼
class EvalInfo:
    def __init__(self, eval_name: str = None, script_name: str = None, source_node: str = None, source_group: str = None, source_id: int = None, source_content_type: str = None, source_content: Union[str, dict] = None, eval_node: str = None, eval_datas: dict[str, list[EvalData]] = None):
        self.eval_name = eval_name
        self.script_name = script_name
        self.source_node = source_node
        self.source_group = source_group
        self.source_id = source_id
        self.source_content_type = source_content_type
        self.source_content = source_content
        self.eval_node = eval_node
        self.eval_datas: dict[str, list[EvalData]] = eval_datas if eval_datas is not None else {}

    def to_dict(self):
        return {
            'eval_name': self.eval_name,
            "script_name": self.script_name,
            "source_node": self.source_node,
            "source_group": self.source_group,
            "source_id": self.source_id,
            "source_content_type":self.source_content_type,
            "source_content": self.source_content,
            "eval_node": self.eval_node,
            "eval_datas": {key: [item.to_dict() for item in val] for key, val in self.eval_datas.items()} if self.eval_datas else {}
        }

    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False,default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()


# 评测结果
class EvalResult:
    def __init__(self, results: dict = None, ext_info: dict = None):
        '''
        :param results: 结果 字典 建议k为data_group
        :param ext_info:扩展数据 字典
        '''
        self.results = results if results else {}
        if not isinstance(self.results, dict):
            try:
                self.results = json.loads(str(self.results))
            except:
                raise ValueError('EvalResult.results must be dict')

        self.ext_info = ext_info if ext_info else {}
        if not isinstance(self.ext_info, dict):
            try:
                self.ext_info = json.loads(str(self.ext_info))
            except:
                raise ValueError('EvalResult.ext_info must be dict')

    def to_dict(self):
        return {
            "results": self.results,
            "ext_info": self.ext_info
        }

    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False, default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()


class EvalTaskCtx():
    def __init__(self, task_name: str = None, eval_name: str = None, req_id: str = None):
        self.task_name = task_name
        self.eval_name = eval_name
        self.req_id = req_id

    def to_dict(self):
        return {
            'task_name': self.task_name,
            'eval_name': self.eval_name,
            'req_id': self.req_id
        }

    # 添加 to_json 方法
    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False, default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()


class PlanKeyDefine:
    def __init__(self, name: str = None, desc: str = None):
        self.name = name
        self.desc = desc

    def to_dict(self):
        return {
            'name': self.name,
            'desc': self.desc
        }

    # 添加 to_json 方法
    def to_json(self):
        return json.dumps(self.to_dict(), ensure_ascii=False, default=str)

    # 实现 __repr__ 方法
    def __repr__(self):
        return self.to_json()

def snode(desc: str = None, content_example: str = None, accept_nodes: list[str] = [], accept_types: list[str] = [], qpm: int = None, rt: int = None, explain_sql: str = 'select * from {table} where {pt};', debug_param_desc: str = None, eval_func: str = None, plan_keys: list[PlanKeyDefine] = None,
          prepare_func: str = None):
    '''
    #  脚本文件随意编写，那么流程怎么串呢？
    # 在你需要将其设置为一个节点的方法加上此装饰器（类似于java注解），系统扫描就会识别到这是一个节点
    # 节点与节点之间可以通过装饰器的accept_nodes属性进行串联，传入值是上游节点的名称，这个accept_nodes属性将单个py脚本里的所有函数串起来
    # 所得链路图为多源多输出的模式，即可以有多个源头，也可以有多个尾节点
    # 每个函数接受两个参数，类型为ScriptNodeData，非空，dict ，非空
    # 返回值可以是NodeResult 也可以是 NodeResult数组 亦或是字典，但是单个元素的结构要和NodeResult相似
    # qpm 以及 rt一起设置时系统会根据配置在跑批时自动设置并发量，支持动态更新
    # accept_types 可限制接受的类型数据，假设我有个节点判断文本是否是饿么诉求，后面还会分别对有/没有的节点跑不同的任务，这时就有用了
    # explain_sql 解释性sql 统一表结构单一无法反映真实节点结构，配置个sql在导出后可以查询想要的结果

    :param desc: 描述 建议
    :param content_example: 接受的模型示例
    :param accept_nodes: 接收的函数列表 根据情况
    :param accept_types: 接受的类型列表 默认情况全要
    :param qpm: 分钟级限流 建议
    :param rt: 预估rt ms 建议
    :param explain_sql: 解释性sql，odps语法 {table} 和 {pt} 分别代表odps表名以及分区 必须存在 例如 select * from {table} where {pt};
    :param debug_param_desc: 调试参数填写说明
    :param eval_func: 在本py存在的函数名 接收俩参数 EvalData 和 EvalTaskCtx
    :param plan_keys: 方案配置的key
    :param prepare_func: worker准备函数
    '''

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 限流
            if qpm is not None:
                rate_limit()

            # 记录开始时间
            start_time = time.time()
            #     统计信息
            monitor_info = __function_monitors.get(func)
            if not monitor_info:
                script_name, node = fetch_key()
                monitor_info = {
                    'details': [],
                    'key': script_name + node,
                    'count_window': 300
                }
                __function_monitors[func] = monitor_info
                __script_schemas.get(script_name).get('functions').get(node)['monitor'] = monitor_info
            try:
                # 调用原始函数
                result = func(*args, **kwargs)
                cost = math.ceil((time.time() - start_time) * 1000)
                monitor_info['details'].append((start_time, cost, None))
                # 10s统计一次
                if monitor_info.get('last_count_time') is None or start_time - monitor_info.get('last_count_time') >= 10:
                    __count_monitor(monitor_info, start_time)
                return result
            except:
                monitor_info['details'].append((start_time, -1, traceback.format_exc()))
                # 失败就计算
                __count_monitor(monitor_info, start_time)
                raise

        def __count_monitor(monitor_info, start_time):
            monitor_info['last_count_time'] = start_time
            # 计算最近5分钟内的调用次数、成功次数和平均耗时
            ten_minutes_ago = start_time - monitor_info['count_window']
            last_error_detail = None
            for sub_arr in reversed(monitor_info['details']):  # 从后往前遍历
                # 如果子元素的第二个元素小于0 就是异常的
                if sub_arr[1] < 0:
                    last_error_detail = sub_arr
                    break
            if last_error_detail:
                monitor_info['last_error_detail'] = last_error_detail
            monitor_info['details'] = [ts for ts in monitor_info['details'] if ts[0] >= ten_minutes_ago]
            recent_call_count = len(monitor_info['details'])
            success_calls = [tup[1] for tup in monitor_info['details'] if tup[1] > 0]
            recent_success_count = len(success_calls)
            recent_avg_time = sum(success_calls) / recent_call_count if recent_call_count > 0 else 0
            util.log('fun_monitor', {
                'fun': monitor_info['key'],
                'recent_call_count': recent_call_count,
                'recent_success_count': recent_success_count,
                'recent_avg_time': recent_avg_time
            })
            # 更新
            monitor_info.update({
                'recent_call_count': recent_call_count,
                'recent_success_count': recent_success_count,
                'recent_avg_time': recent_avg_time
            })

        # 限流
        def rate_limit():
            if func not in __function_buckets:
                rate_bucket = RateBucket(qpm=qpm)
                __function_buckets[func] = rate_bucket
            bucket = __function_buckets[func]
            # 每次等待限流的1/10
            once_wait = 60 / qpm / 10
            wait_time = 0
            while True:
                if bucket.try_fill():
                    break
                else:
                    if wait_time >= 60:
                        raise Exception("Timeout waiting for function call")
                    time.sleep(once_wait)
                    wait_time += once_wait

        def fetch_key():
            script_name = sys.modules[func.__module__].__file__
            if script_directory not in script_name:
                raise ValueError('脚本必须在指定目录下,' + script_directory)
            script_name = script_name.replace(script_directory + '/', "", 1).split('.')[0]
            node = func.__name__
            return script_name, node

        return wrapper

    return decorator


# 漏桶限流
class RateBucket:
    def __init__(self, qpm):
        self.capacity = max(math.floor(qpm / 60.0), 1)
        self.water = 0
        self.last_leak_time = time.time()
        self.leak_rate = qpm / 60.0

    def leak_out(self):
        now = time.time()
        time_since_last_leak = now - self.last_leak_time
        leaked_water = time_since_last_leak * self.leak_rate
        self.water = max(self.water - leaked_water, 0)
        self.last_leak_time = now

    def try_fill(self):
        self.leak_out()
        if self.water < self.capacity:
            self.water += 1
            return True
        else:
            return False


if sys.version_info < (3, 9):
    import astor

    unparse = astor.to_source
else:
    unparse = ast.unparse


class SNodeAnnotationAnalyzer(ast.NodeVisitor):
    def __init__(self, script_key, script_name):
        self.dependencies = set()
        self.functions = {}
        self.script_key = script_key
        self.script_name = script_name

    def visit_Import(self, node: ast.Import):
        for alias in node.names:
            self.dependencies.add(alias.name)
        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef):
        desc = node.name
        accept_nodes = []
        accept_types = []
        content_example = None
        qpm = None
        rt = None
        explain_sql = None
        suggest_worker_num = None
        eval_func = None
        prepare_func = None
        has_snode = False
        debug_param_desc = None
        plan_keys = []
        for dec in node.decorator_list:
            if isinstance(dec, ast.Call) and dec.func.id == 'snode':
                has_snode = True
                for kw in dec.keywords:
                    if kw.arg == 'desc':
                        desc = kw.value.s
                    elif kw.arg == 'accept_nodes':
                        accept_nodes = [arg.s for arg in kw.value.elts]
                    elif kw.arg == 'accept_types':
                        accept_types = [arg.s for arg in kw.value.elts]
                    elif kw.arg == 'qpm':
                        qpm = int(kw.value.s)
                    elif kw.arg == 'rt':
                        rt = int(kw.value.s)
                    elif kw.arg == 'explain_sql':
                        explain_sql = kw.value.s
                        if '{table}' not in explain_sql or '{pt}' not in explain_sql:
                            raise ValueError('explain_sql must has {table} and {pt},now the explain_sql is ' + explain_sql)
                        # 替换实体
                        if one_runtime.get_config('WRITE_ODPS_PROJECT'):
                            explain_sql = explain_sql.replace('{table}', one_runtime.get_config('WRITE_ODPS_PROJECT') + '.' + one_runtime.get_config('WRITE_ODPS_TABLE'))
                            explain_sql = explain_sql.replace('{pt}', 'pt=\'${pt}\'')
                    elif kw.arg == 'content_example':
                        content_example = kw.value.s
                    elif kw.arg == 'debug_param':
                        debug_param_desc = kw.value.s
                    elif kw.arg == 'eval_func':
                        eval_func = kw.value.s
                    elif kw.arg == 'prepare_func':
                        prepare_func = kw.value.s
                    elif kw.arg == 'plan_keys':
                        for item in kw.value.elts:
                            if isinstance(item, ast.Call):
                                # 提取位置参数
                                positional_args = [ast.literal_eval(a) for a in item.args]

                                # 提取关键字参数
                                keyword_args = {kw.arg: ast.literal_eval(kw.value) for kw in item.keywords}

                                # 使用位置参数和关键字参数创建自定义对象
                                # 这里我们假设你有一个PlanKeyDefine类，并且它的__init__方法接受这些参数
                                plan_keys.append(PlanKeyDefine(*positional_args, **keyword_args).to_dict())

        if not has_snode:
            return
        if not re.compile('^[a-zA-Z_0-9]+$').match(node.name):
            raise ValueError(f"Function {node.name} must be [a-zA-Z_0-9]+")

        if node.name == self.script_name:
            raise ValueError('节点名和脚本名不可相同')

        if qpm is not None and rt is not None:
            # 最小为1
            suggest_worker_num = math.ceil(qpm / (60 / (rt / 1000.0)))

        if len(node.args.args) != 2:
            raise ValueError(f"Function {node.name} must have exactly two arguments.")

        # Insert the function name and arguments into the function body
        arg1 = node.args.args[0].arg
        arg2 = node.args.args[1].arg
        func_signature = f"def {node.name}({arg1}:ScriptNodeData,{arg2}:NodeTaskCtx):"

        # Extract the function body
        body = unparse(node.body).strip()
        # Add a tab to the beginning of each line in the function body
        body = '\n'.join(['\t' + line for line in body.split('\n')])
        # 使用正则表达式匹配字符串中的 '''xxx\nxxx'''
        # 使用正则表达式匹配字符串中的 'xxx\nxxx'
        pattern = re.compile(r"'([^']*\\n[^']*)'")
        # 对匹配到的每个字符串进行处理
        body = pattern.sub(lambda m: "'''" + m.group(1).replace('\\n', '\n') + "'''", body)
        # Combine function signature with the body
        full_function = func_signature + '\n' + body

        self.functions[node.name] = {
            'accept_nodes': accept_nodes,
            'root': len(accept_nodes) == 0,
            'accept_types': accept_types,
            'body': full_function,
            'desc': desc,
            'qpm': qpm,
            'rt': rt,
            'suggest_worker_num': suggest_worker_num,
            'explain_sql': explain_sql,
            'content_example': content_example,
            'debug_param_desc': debug_param_desc,
            'eval_func': eval_func,
            'plan_keys': plan_keys,
            'prepare_func': prepare_func
        }
        self.generic_visit(node)


def analyze_code(script_key, script_name, code: str):
    tree = ast.parse(code)
    analyzer = SNodeAnnotationAnalyzer(script_key, script_name)
    analyzer.visit(tree)
    return {
        'dependencies': list(analyzer.dependencies),
        'functions': analyzer.functions
    }


def __dfs(node, graph, visited, recursion_stack, path):
    visited[node] = True
    recursion_stack[node] = True
    path.append(node)

    for neighbor in graph[node]['accept_nodes']:
        if not visited.get(neighbor, False):
            __dfs(neighbor, graph, visited, recursion_stack, path)
        elif recursion_stack.get(neighbor, False):
            raise ValueError('Cycle:', ' -> '.join(path + [neighbor]))

    path.pop()
    recursion_stack[node] = False


def is_cyclic(graph):
    visited = {}
    recursion_stack = {}
    path = []

    for node in graph.keys():
        if not visited.get(node, False):
            __dfs(node, graph, visited, recursion_stack, path)


def find_previous_nodes(node, previous_nodes, graph):
    for neighbor in graph[node]['accept_nodes']:
        if node not in previous_nodes[neighbor]:
            previous_nodes[neighbor].append(node)
            find_previous_nodes(neighbor, previous_nodes, graph)


def update_graph_reverse(graph):
    previous_nodes = {node: [] for node in graph.keys()}
    for node in graph.keys():
        find_previous_nodes(node, previous_nodes, graph)
    for node in graph.keys():
        graph[node]['next_nodes'] = previous_nodes[node]


def process_script_code(script_name, python_code: str) -> dict:
    pinyin_list = pinyin(script_name, style=Style.NORMAL, strict=True)
    script_key = ''.join([item[0] if not item[0].isalpha() else item[0].lower() for item in pinyin_list])
    analysis = analyze_code(script_key, script_name, python_code)
    dependencies = analysis['dependencies']
    functions = analysis['functions']
    # 环检测
    is_cyclic(functions)
    # 下游验证
    update_graph_reverse(functions)
    return {
        'script_name': script_name,
        'script_key': script_key,
        'dependencies': dependencies,
        'functions': functions
    }


def scan_for_py_files(path):
    py_files = []
    for root, dirs, files in os.walk(path):
        for file in files:
            if file.endswith(".py"):
                py_files.append(os.path.join(root, file))
    return py_files


def parse_script_file(file_path):
    file_name = os.path.basename(file_path).split('.')[0]
    if file_name == '__init__':
        return None
    with open(file_path, "r") as file:
        content = file.read()
        return process_script_code(file_name, content)


def refresh_scripts():
    global __script_schemas, __script_modules, __function_monitors, __function_buckets
    temp = {}
    temp_modules = {}
    try:
        util.log('begin_refresh_scripts')
        if not os.path.isdir(script_directory):
            util.log('script_init', {}, "The 'script' folder does not exist in the current directory.")
            return

        py_files = scan_for_py_files(script_directory)
        for py_file in py_files:
            script = parse_script_file(py_file)
            if not script or not script.get('functions'):
                continue
            script_key_ = script['script_key']
            if script_key_ in temp:
                raise ValueError('同在同key脚本', script_key_)
            script_name = script['script_name']
            # 使用字符串的 replace 方法，将 a 中的 b 替换为空字符串
            module_name = py_file.replace(script_directory, "", 1)

            # 移除 '.py' 后缀
            if module_name.endswith('.py'):
                module_name = module_name[:-3]
            module_name = 'script.scripts' + module_name.replace('/', ".")
            script.update({
                'module_name': module_name
            })
            # 刷新
            try:
                module = importlib.import_module(module_name)
                temp_modules.update({
                    script_name: module
                })
                importlib.reload(module)

                for key, val in script.get('functions').items():
                    eval_func = val.get('eval_func')
                    if eval_func:
                        # 解析 & 验证
                        func = getattr(module, eval_func)
                        if not func:
                            raise ValueError('eval function is not exist')
                        # 检查函数参数数量
                        sig = inspect.signature(func)
                        parameters = sig.parameters
                        if len(parameters) != 2:
                            raise ValueError('eval function must have two parameters')

                    prepare_func = val.get('prepare_func')
                    if prepare_func:
                        # 解析 & 验证
                        func = getattr(module, prepare_func)
                        if not func:
                            raise ValueError('prepare function is not exist')
                        # 检查函数参数数量
                        sig = inspect.signature(func)
                        parameters = sig.parameters
                        if len(parameters) != 2:
                            raise ValueError('prepare function must have two parameters')

                temp.update({
                    script_name: script
                })
            except Exception as e:
                util.log('module_update_error', {
                    'error': e
                }, f"Error importing module {module_name}")  # 动态导入 失败
    finally:
        __script_schemas = temp
        __script_modules = temp_modules
        __function_monitors = {}
        __function_buckets = {}
        util.log('end_refresh_scripts')

